import multiprocessing

class HmmSearchParser( ):
    def __init__(self, lock=None, verbose=False, domtblout=False):
        self.lock = lock if lock is None else multiprocessing.Lock()
        if lock is None:
            self.lock = multiprocessing.Lock()
        else:
            self.lock = lock
        if verbose:
            self.print_headers = True
            self.verbose = True
        if domtblout:
            self.set_scorer(True)
        else:
            self.set_scorer(False)
        self.line_splitter = re.compile('\s+')

    ## Line parsing
    def _parse_line(self, *args):
        """API for attaching line parsing methods. Have verbose & non-verbose,
        configured with ``set_verbose( True|False )''
        """
        pass

    def _parse_verbose(self, line):
        with self.lock:
            print(line),
        return self.get_scores( line )

    def _silent_parse( self, line ):
        return self.get_scores( line )

    def set_verbose(self, print_line=False):
        if print_line:
            self.verbose = True
            self._parse_line = self._parse_verbose
        else:
            self.verbose = False
            self._parse_line = self._silent_parse

    ## Extracting scores
    def get_scores(self, *args):
        """API for attaching hmmsearch result parsers. Currently have a table
        and domain table parser, tested against hmmsearch v3.0.
        Can switch between the two with set_scorer( True|False ).
        """

    def _domtbl_scores(self, line):
        values = self.line_splitter.split(line.rstrip(), 22)
        #                                                                                     --- full sequence --- -------------- this domain -------------           hmm    coord    ali     coord   env coord
        # target name                 accession   tlen query name           accession   qlen   E-value  score  bias   #      of    c-Evalue  i-Evalue  score    bias  from       to    from       to      from    to     acc     description of target
        # accession,                  null,      seqLen,  HMM,                null2,   HMMLen,   Eval,  score, bias,nDomains,null3,cEval,   iEval,    domscore,dombias,hmmstart,hmmend,alistart,aliend,envstart,envend,accuracy,description
        #  0                          1           2        3                    4       5         6      7     8      9      10    11       12         13        14    15        16    17         18       19     20      21       22
        result_dict = {'accession'  : values[0]        ,  # accession,
                       'seqLen'     : int(values[2])   ,  # seqLen,
                       'HMM'        : values[3]        ,  # HMM,
                       'HMMLen'     : values[5]        ,  # HMMLen,
                       'Eval'       : float(values[6]) ,  # Eval,
                       'score'      : float(values[7]) ,  # score,
                       'bias'       : float(values[8]) ,  # bias,
                       '#'          : int(values[9])   ,  # nDomains,
                       'cEval'      : float(values[11]),  # cEval,
                       'iEval'      : float(values[12]),  # iEval,
                       'domscore'   : float(values[13]),  # domscore,
                       'dombias'    : float(values[14]),  # dombias,
                       'hmmstart'   : int(values[15])  ,  # hmmstart,
                       'hmmend'     : int(values[16])  ,  # hmmend,
                       'alistart'   : int(values[17])  ,  # alistart,
                       'aliend'     : int(values[18])  ,  # aliend,
                       'envstart'   : int(values[19])  ,  # envstart,
                       'envend'     : int(values[20])  ,  # envend,
                       'accuracy'   : float(values[21]),  # accuracy,
                       'description': values[22],  # seqDesc,
                       }
        return result_dict

    def _tbl_scores(self, line):
        values = self.line_splitter.split( line )
        result_dict = { 'accession' : values[0],
                        'Eval'      : float(values[4]),
                        'score'     : float(values[5]),
                }
        return result_dict

    def set_scorer(self, domain_table):
        """If parsing hmmsearch results from the `--domtblout` option,
        then the tabular output is different. Pass True to parse
        domain table results. This is the default.
        Otherwise, pass False
        """
        if domain_table:
            self.get_scores = self._domtbl_scores
        else:
            self.get_scores = self._tbl_scores

    ## Header parsing
    def _parse_header(self, *args):
        """API for parsing headers. If you want to change from silent
        to verbose parsing or vice-versa, use `set_parse_header(bool)`"""

    def _silent_header(self, line):
        return

    def _parse_header_verbose(self, line):
        with self.lock:
            print(line)
        return

    def set_parse_header(self, print_header=True):
        if print_header:
            self.print_headers = True
            self._parse_header = self._parse_header_verbose
        else:
            self.print_headers = False
            self._parse_header = self._silent_header

    def parse(self, hmmsearchProcess):
        ## *** The parser!! *** ##
        try:
            for line in hmmsearchProcess.stdout:
                if line[:1] == '#':
                    self._parse_header( line )
                else:
                    yield self._parse_line( line )
            retCode = hmmsearchProcess.wait()
            if retCode != 0:
                sys.stderr.write( hmmsearchProcess.stderr.read() )
                sys.stderr.flush()
                raise IOError
        except KeyboardInterrupt:
            pass
        except Exception:
            raise
        finally:
            hmmsearchProcess.stdout.close()
            hmmsearchProcess.stderr.close()


#class HmmSearcher( Thread):
        #Thread.__init__(self)
class HmmSearcher(multiprocessing.Process):
    """
    """
    def __init__(self, inPipe, outPipe, d_Q, lock):
        multiprocessing.Process.__init__(self)
        if os.uname()[-1] == 'x86_64':
            self.hmmsearch = os.path.join(CONFIG.hmmerdir, CONFIG.hmmsearchCMD)
        else:
            self.hmmsearch = os.path.join(CONFIG.hmmer32dir, CONFIG.hmmsearchCMD)
        self.inPipe, self.__outPipe = (inPipe, outPipe, )
        self.outQ = d_Q   # Goes to Scorer.
        #self.sem = semaphore
        if CONFIG.options['--max']: # max sensitivity:-
            self.command = [ self.hmmsearch, '-o', os.devnull, '--cpu', '1',
                             '--max',
                             '--domtblout', '/dev/stdout', '--noali',
                             'REPLACED', '/dev/stdin' ]
        else:
            self.command = [ self.hmmsearch, '-o', os.devnull, '--cpu', '1',
                             '--domtblout', '/dev/stdout', '--noali',
                             'REPLACED', '/dev/stdin' ]
        CONFIG.printHeaders = False
        self.lock = lock
        self.parser = HmmSearchParser()
        self.parser.set_scorer(True)  # parse domtblout instead of tblout

    def run(self):
        self.Eval = CONFIG.options['-Eval']
        self.score = CONFIG.options['-score']
        parser = self.parser.parse
        if CONFIG.options['--verbose']:
            self.parser.set_verbose(True)
        else:
            self.parser.set_verbose(False)
        try:
            OTUName = None
            nseqs, HMMLocation = self.__outPipe.recv()  # <-- SequenceDistributor
            while nseqs != 'END':
                self.command[-2] = HMMLocation  ### This replaces 'REPLACED' in self.command, setting the HMM location.
                OTUName = HMMLocation[ HMMLocation.rfind( os.path.sep ) + 1: HMMLocation.rfind('.') ]
                try:
                    process = subprocess.Popen(self.command, shell=False,
                                               stdin=subprocess.PIPE,
                                               stderr=subprocess.PIPE,
                                               stdout=subprocess.PIPE,
                                               bufsize=-1)
                except OSError, e:
                    self.outQ.put( ('die', e, OTUName )) # --> Scorer
                    sys.stderr.write( "COMMAND FAILED: {0}\n".format( ' '.join(self.command) ) )
           #         while self.__outPipe.recv() != 'END':  # <-- should come from distributor
           #             continue
                    self.__outPipe.recv() # seqs
                    nseqs, HMMLocation = self.__outPipe.recv()
                    continue
                try:
                    # feed in the sequences received from distributor.
                    for seq in self.__outPipe.recv():
                        process.stdin.write( seq )
                except IOError, e:
                    #sys.stderr.write('[HmmSearcher.run] {0} dying. --> Scorer'.format( OTUName ) )
                    hmmerror = '\n'.join( process.communicate() )
                    self.outQ.put( ('die', '{0}\n{1}'.format(e, hmmerror), OTUName )) # --> Scorer
                    sys.stdout.write('HMMError: ')
                    sys.stdout.write(hmmerror)
           #         while self.__outPipe.recv() != 'END':  # <-- should come from distributor
           #             continue
                    raise(e)
                except Exception, e:
                    hmmerror = '\n'.join( process.communicate() )
                    hmmerror = process.communicate()[1]
                    self.outQ.put( ('die', '{0}\n{1}'.format(e, hmmerror), OTUName) ) # --> Scorer
                    raise(e)
                finally:
                    process.stdin.close()
                for results in parser( process ):
                    self.outQ.put( ('update', OTUName, results, )  )   # --> Scorer
                #self.sem.release()  # Place if using semaphores, which seem unnecessary..?
                self.outQ.put( ('end', OTUName, ) )        ## --> Scorer
                nseqs, HMMLocation = self.__outPipe.recv() ## <--  distributor
        except KeyboardInterrupt, e:
            self.outQ.put( ('die', e, {'OTUName' : OTUName}) ) # --> Scorer
            raise(e)
        finally:
            self.__outPipe.close()
        return

